#!/usr/bin/python
# Filename: tag.py

import math;

def max_dict(s):
	value = float('-inf');
	key = '';
	for name in s:
		if s[name] >= value:
			key = name;
			value = s[name];
	return (key,value);

def Viterbi(s,q,e,tags):
	pi = {};
	bp = {};
	pi_dict = {};
	tag_prob = {};
	
	x = s;
	
	pi[0,'*','*'] = math.log(float(1));#initialization
	n = len(x);
	tags.add('*');
	for tag_1 in tags:
		for tag_2 in tags:
			if tag_1 != '*' or tag_2 != '*':
				pi[0,tag_1,tag_2] = float('-inf');

	for k in range(1,n+1):
		if not(x[k] in e):
			x[k] = '_RARE_';
		for tag_u in tags:
			for tag_v in tags:
				pi_dict[tag_u,tag_v] = {};
				for tag_w in tags:
					if not((tag_w,tag_u,tag_v) in q):
						q[tag_w,tag_u,tag_v] = float('-inf');
					if not(tag_v in e[x[k]]):
						e[x[k]][tag_v] = float('-inf');
					pi_dict[tag_u,tag_v][tag_w]= pi[k-1,tag_w,tag_u]+q[tag_w,tag_u,tag_v]+e[x[k]][tag_v];
				pi_collect = {};
				pi_collect = pi_dict[tag_u,tag_v];
				(w,v) = max_dict(pi_collect);
				pi[k,tag_u,tag_v] = v;
				bp[k,tag_u,tag_v] = w;

	tag_tbl = {};
	tag={}
	for tag_u in tags:
		for tag_v in tags:
			if not((tag_u,tag_v,'STOP') in q):
				q[tag_u,tag_v,'STOP'] = float('-inf');
			tag_tbl[tag_u,tag_v] = pi[n,tag_u,tag_v]+q[tag_u,tag_v,'STOP'];

	name,v = max_dict(tag_tbl);

	tag[n-1] = name[0];
	tag[n] = name[1];
	for k in range(n-2,0,-1):
		tag[k] = bp[k+2,tag[k+1],tag[k+2]];
	tag[0] = '*';
	for k in range(1,n+1):
		tag_prob[k] = pi[k,tag[k-1],tag[k]];
	
	return (tag,tag_prob)
	


q_tbl = dict();
tags = set();
e_tbl = dict();

fr_e_dat = file('emission_rep.dat');
for line in fr_e_dat:    #read e(x|y) into e_tbl and add all the tags into tags
	words = line.strip().split(' ');
	if words[0] in e_tbl:
		e_tbl[words[0]][words[1]] = float(words[2]);
	else:
		e_tbl[words[0]] = {words[1] : float(words[2])}; 
	if not(tags.__contains__(words[1])):
			tags.add(words[1]);
fr_e_dat.close();

fr_q_dat = file('trigram_prob.dat');
for line in fr_q_dat: 	#read q(s|u,v) into q_tbl
	words = line.strip().split(' ');
	q_tbl[words[0],words[1],words[2]] = float(words[3]);
fr_q_dat.close();


fr_dev = file('ner_dev.dat');

fw_dev = open('ner_dev_viterbi.tag','w');

sentence = {};
i = 1; 
for line in fr_dev:
	words = line.strip().split(' ');
	if(words[0] != ''):
		sentence[i] = words[0];
		i += 1;
		continue;
	else:
		x = sentence.copy();
		tag,tag_prob = Viterbi(x,q_tbl,e_tbl,tags);
		for j in range(1,len(sentence)+1):
			lines = sentence[j] + ' ' + tag[j] + ' ' + str(tag_prob[j]) + '\n';
			fw_dev.write(lines);
		i = 1;
		fw_dev.write('\n');
		sentence.clear();
fw_dev.close();
fr_dev.close();
		
